#include <jni.h>
#include <string>
#include "libyuv.h"

void I420ToNV21(const uint8_t* i420_data, uint8_t*  nv21_data, jint width, jint height, bool isI420);
void I420Rotate(const uint8_t* src, uint8_t* dst, jint width, jint height, jint rotation);

extern "C"
JNIEXPORT void JNICALL
Java_cn_android666_libyuv_Libyuv_yv12ToI420(JNIEnv *env, jclass clazz, jbyteArray yv12_bytes,
                                            jbyteArray i420_bytes, jint width, jint height) {
    // 第二个参数是 jboolean* isCopy，它的作用是告诉你 JNI 是否返回了原始数组的直接指针，还是创建了一个本地拷贝。
    // 如果*isCopy == JNI_TRUE，则JNI返回了一个本地拷贝
    // 如果*isCopy == JNI_FALSE，则JNI返回了Java数组的直接指针（没有拷贝）。
    // 如果你不关心 JNI 是否创建了拷贝，只想拿到一个可以访问数据的指针，则可以传nullptr。
    // GetByteArrayElements() 到底是直接返回java数组的内存地址还是返回一个拷贝，这是由JVM决定的，我们控制不了。
    jbyte* src = env->GetByteArrayElements(yv12_bytes, nullptr);
    jbyte* dst = env->GetByteArrayElements(i420_bytes, nullptr);
    if (!src || !dst) return;

    int y_size = width * height;
    int uv_size = (width / 2) * (height / 2);

    const uint8_t* src_y = reinterpret_cast<uint8_t*>(src);
    const uint8_t* src_v = src_y + y_size;
    const uint8_t* src_u = src_v + uv_size;

    auto* dst_y = reinterpret_cast<uint8_t*>(dst);
    uint8_t* dst_u = dst_y + y_size;
    uint8_t* dst_v = dst_u + uv_size;

    // 直接拷贝到 I420
    libyuv::I420Copy(src_y, width,
                     src_u, width / 2,
                     src_v, width / 2,
                     dst_y, width,
                     dst_u, width / 2,
                     dst_v, width / 2,
                     width, height);

    // ReleaseByteArrayElements参数介绍：
    // 参数1：要释放的数组对象
    // 参数2：要释放的数组指针
    // 参数3：告诉 JVM 是否需要把本地修改写回 Java 数组。
    //      0：写回并释放本地指针。修改会同步到 Java 数组。
    //      JNI_COMMIT: 写回，但不释放本地指针（还可以继续使用）。
    //      JNI_ABORT：不写回，直接释放本地指针。修改丢弃。在这个示例中也表示src只是用来读取yv12_bytes，没有修改，所以也不需要写回，避免不必要的开销。
    env->ReleaseByteArrayElements(yv12_bytes, src, JNI_ABORT); // JNI_ABORT: 丢弃修改，直接释放本地指针，不写回 Java 数组
    env->ReleaseByteArrayElements(i420_bytes, dst, 0);         // 0：将修改写回 Java 数组（如果使用了拷贝）并释放本地指针
}

extern "C"
JNIEXPORT void JNICALL
Java_cn_android666_libyuv_Libyuv_yv12ToNv21(JNIEnv *env, jclass clazz, jbyteArray yv12_bytes,
                                            jbyteArray nv21_bytes, jint width, jint height) {
    jbyte* yv12_data = env->GetByteArrayElements(yv12_bytes, nullptr);
    jbyte* nv21_data = env->GetByteArrayElements(nv21_bytes, nullptr);

    I420ToNV21(reinterpret_cast<uint8_t*>(yv12_data), reinterpret_cast<uint8_t*>(nv21_data), width, height, false);

    env->ReleaseByteArrayElements(yv12_bytes, yv12_data, JNI_ABORT);
    env->ReleaseByteArrayElements(nv21_bytes, nv21_data, 0);
}

void I420ToNV21(const uint8_t* i420_data, uint8_t* nv21_data, jint width, jint height, bool isI420) {
    int y_size = width * height;
    int uv_size = (width / 2) * (height / 2);

    const uint8_t* src_y = i420_data;
    const uint8_t* src_u;
    const uint8_t* src_v;

    if (isI420) {
        // I420: Y - U - V
        src_u = i420_data + y_size;
        src_v = i420_data + y_size + uv_size;
    } else {
        // YV12: Y - V - U
        src_v = i420_data + y_size;
        src_u = i420_data + y_size + uv_size;
    }

    uint8_t* dst_y = nv21_data;
    uint8_t* dst_vu = nv21_data + y_size;

    int src_stride_y = width;
    int src_stride_u = width / 2;
    int src_stride_v = width / 2;
    int dst_stride_y = width;
    int dst_stride_vu = width;

    libyuv::I420ToNV21(src_y, src_stride_y,
               src_u, src_stride_u,
               src_v, src_stride_v,
               dst_y, dst_stride_y,
               dst_vu, dst_stride_vu,
               width, height);
}

void I420Rotate(const uint8_t* src, uint8_t* dst, jint width, jint height, jint rotation) {
    int y_size = width * height;
    int uv_width = width / 2;
    int uv_height = height / 2;
    int uv_size = uv_width * uv_height;

    const uint8_t* src_y = src;
    const uint8_t* src_v = src + y_size;
    const uint8_t* src_u = src + y_size + uv_size;

    int dst_width = width;
    int dst_height = height;

    libyuv::RotationMode mode;
    switch (rotation) {
        case 0:   mode = libyuv::kRotate0;   break;
        case 90:  mode = libyuv::kRotate90;  dst_width = height; dst_height = width; break;
        case 180: mode = libyuv::kRotate180; break;
        case 270: mode = libyuv::kRotate270; dst_width = height; dst_height = width; break;
        default:  mode = libyuv::kRotate0;   break;
    }

    // 输出的UV平面大小要根据旋转后的宽高重新计算
    int dst_uv_width = dst_width / 2;
    int dst_uv_height = dst_height / 2;
    int dst_y_size = dst_width * dst_height;
    int dst_uv_size = dst_uv_width * dst_uv_height;

    uint8_t* dst_y = dst;
    uint8_t* dst_v = dst_y + dst_y_size;
    uint8_t* dst_u = dst_v + dst_uv_size;

    libyuv::I420Rotate(
            src_y, width,
            src_u, uv_width,
            src_v, uv_width,
            dst_y, dst_width,
            dst_u, dst_uv_width,
            dst_v, dst_uv_width,
            width, height,
            mode
    );
}

extern "C"
JNIEXPORT void JNICALL
Java_cn_android666_libyuv_Libyuv_nv21ToNv12(JNIEnv *env, jclass clazz, jbyteArray nv21_bytes,
                                            jbyteArray nv12_bytes, jint width, jint height) {
    jbyte* src = env->GetByteArrayElements(nv21_bytes, nullptr);
    jbyte* dst = env->GetByteArrayElements(nv12_bytes, nullptr);
    if (!src || !dst) return;

    int y_size = width * height;

    const uint8_t* src_y = reinterpret_cast<uint8_t*>(src);
    const uint8_t* src_vu = reinterpret_cast<uint8_t*>(src + y_size);

    auto* dst_y = reinterpret_cast<uint8_t*>(dst);
    auto* dst_uv = reinterpret_cast<uint8_t*>(dst + y_size);

    libyuv::NV21ToNV12(src_y, width,
                       src_vu, width,
                       dst_y, width,
                       dst_uv, width,
                       width, height);

    env->ReleaseByteArrayElements(nv21_bytes, src, JNI_ABORT);
    env->ReleaseByteArrayElements(nv12_bytes, dst, 0);
}

extern "C"
JNIEXPORT void JNICALL
Java_cn_android666_libyuv_Libyuv_yv12Rotate(JNIEnv *env, jclass clazz, jbyteArray yv12_bytes,
                                            jbyteArray result_bytes, jint width, jint height,
                                            jint rotation) {
    jbyte* src = env->GetByteArrayElements(yv12_bytes, nullptr);
    jbyte* dst = env->GetByteArrayElements(result_bytes, nullptr);
    if (!src || !dst) return;

    I420Rotate(reinterpret_cast<const uint8_t*>(src), reinterpret_cast<uint8_t*>(dst), width, height, rotation);

    env->ReleaseByteArrayElements(yv12_bytes, src, JNI_ABORT);
    env->ReleaseByteArrayElements(result_bytes, dst, 0);
}

extern "C"
JNIEXPORT void JNICALL
Java_cn_android666_libyuv_Libyuv_yv12RotateToNv21(JNIEnv *env, jclass clazz, jbyteArray yv12_bytes,
                                                  jbyteArray result_bytes, jbyteArray temp_bytes, jint width, jint height,
                                                  jint rotation) {
    jbyte* src = env->GetByteArrayElements(yv12_bytes, nullptr);
    jbyte* dst = env->GetByteArrayElements(result_bytes, nullptr);
    jbyte* temp = env->GetByteArrayElements(temp_bytes, nullptr);
    if (!src || !dst || !temp) return;

    I420Rotate(reinterpret_cast<const uint8_t*>(src), reinterpret_cast<uint8_t*>(temp), width, height, rotation);
    I420ToNV21(reinterpret_cast<const uint8_t*>(temp), reinterpret_cast<uint8_t*>(dst), width, height, false);

    env->ReleaseByteArrayElements(yv12_bytes, src, JNI_ABORT);
    env->ReleaseByteArrayElements(result_bytes, dst, 0);
    env->ReleaseByteArrayElements(temp_bytes, temp, JNI_ABORT);
}